from os import listdir
from os.path import isfile, join
import tensorflow as tf


def get_image(path, height, width, preprocess_fn):
    png = path.lower().endswith('png')
    img_bytes = tf.read_file(path)
    image = tf.image.decode_png(img_bytes, channels=3) if png else tf.image.decode_jpeg(img_bytes, channels=3)
    return preprocess_fn(image, height, width)

'''image(batch_size, height, width, path, preprocess_fn, epochs=2, shuffle=True)
    batch_size:批处理大小
    height, width：图像的高和宽
    path：训练集的path
    preprocess_fn：预处理返回的结果tf.image.per_image_whitening(resized_image)
    shuffle：打乱顺序
    '''
def image(batch_size, height, width, path, preprocess_fn, epochs=2, shuffle=True):
    '''遍历数据集，join()拼接文件名,返回一个列表'''
    filenames = [join(path, f) for f in listdir(path) if isfile(join(path, f))]
    if not shuffle:
        filenames = sorted(filenames)
    '''文件名变小写，如果第一个文件是PNG，则假定它们都是 '''
    png = filenames[0].lower().endswith('png')
    '''
string_input_producer(
    string_tensor,为提供的文件列表
    num_epochs=None, 用于限制加载出事文件列表的最大轮数,当设置了此参数后，会在本地计数，加载次数结束后会报 OutOfRange 错误
    shuffle=True,
    seed=None,
    capacity=32,为队列的容量
    shared_name=None,
    name=None,
    cancel_op=None
)返回一个带有输出字符串的队列
'''
    filename_queue = tf.train.string_input_producer(filenames, shuffle=shuffle, num_epochs=epochs)
    reader = tf.WholeFileReader()#一个阅读器,读取整个文件,返回文件名称key,以及文件中所有的内容value
    _, img_bytes = reader.read(filename_queue)#返回A tuple of Tensors (key, value).
    image = tf.image.decode_png(img_bytes, channels=3) if png else tf.image.decode_jpeg(img_bytes, channels=3)#RGB三通道
    '''预处理图片，返回处理过的图片'''
    processed_image = preprocess_fn(image, height, width)
    '''tf.train.batch是按顺序读取数据'''
    return tf.train.batch([processed_image], batch_size, dynamic_pad=True)
